import os
import shutil
from pathlib import Path

def sync_dataset():
    # Define paths
    base_dir = Path('dataset')
    vis_dir = Path('visualization')
    images_dir = base_dir / 'images'
    labels_dir = base_dir / 'labels'
    
    # Create directories if they don't exist
    for split in ['train', 'val']:
        (images_dir / split).mkdir(parents=True, exist_ok=True)
        (labels_dir / split).mkdir(parents=True, exist_ok=True)
    
    # Get list of images from visualization folder
    vis_images = set(f.name for f in vis_dir.glob('*.jpg'))
    
    # Process each split
    for split in ['train', 'val']:
        # Get current images in the split
        current_images = set(f.name for f in (images_dir / split).glob('*.jpg'))
        
        # Find images to keep (those in visualization)
        images_to_keep = current_images.intersection(vis_images)
        
        # Remove images not in visualization
        for img in current_images - vis_images:
            img_path = images_dir / split / img
            label_path = labels_dir / split / img.replace('.jpg', '.txt')
            
            if img_path.exists():
                img_path.unlink()
            if label_path.exists():
                label_path.unlink()
        
        print(f"Split {split}:")
        print(f"Total images: {len(current_images)}")
        print(f"Images kept: {len(images_to_keep)}")
        print(f"Images removed: {len(current_images - vis_images)}")
        print()

if __name__ == '__main__':
    sync_dataset() 